package com.wulin.wucut.utils;

import io.github.mymonstercat.Model;
import io.github.mymonstercat.ocr.InferenceEngine;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

@Slf4j
public class ORCUtil {
    
    private ORCUtil(){}
    private static final InferenceEngine engine = InferenceEngine.getInstance(Model.ONNX_PPOCR_V3);

    public static String doOcr(String image) {
        return engine.runOcr(image).getStrRes();
    }

    public static List<File> getFramesWritePic(String videoPath) throws IOException {
        List<File> picFile = new ArrayList<>();
        FFmpegFrameGrabber frameGrabber = new FFmpegFrameGrabber(videoPath.replaceAll("\\\\","/"));
        frameGrabber.start();
        long lengthInTime = frameGrabber.getLengthInTime();
        int step =  (int) (0.5 * 1000 * 1000);
        long framesIndex = 1 * 1000 * 1000;
        int index = 0;
        List<BufferedImage> bufferedImageList = new ArrayList<>();
        while (framesIndex <= lengthInTime) {
            frameGrabber.setTimestamp(framesIndex);
            Frame frame = frameGrabber.grabImage();
            if (Objects.nonNull(frame)) {
                BufferedImage bufferedImage = frameToBufferedImage(frame);
                bufferedImageList.add(bufferedImage);
            }
            framesIndex += step;
            index++;
            log.info("第{}次切割",index);
        }
        for (int i = 0; i < bufferedImageList.size(); i++) {
            //写入输出流
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            ImageIO.write(bufferedImageList.get(i), "png", outputStream);
            File tempFile = File.createTempFile("image" + i, ".png");
            FileOutputStream fos = new FileOutputStream(tempFile);
            fos.write(outputStream.toByteArray());
            picFile.add(tempFile);
        }
        return picFile;
        
    }
    public static BufferedImage frameToBufferedImage(Frame frame) {
        //创建BufferedImage对象
        Java2DFrameConverter converter = new Java2DFrameConverter();
        BufferedImage bufferedImage = converter.getBufferedImage(frame);
        return bufferedImage;
    }

    public static void main(String[] args) {
        String s = "“你想和解吗？";
        System.out.println(s);
    }
}
